#include "common.h"

// 下载状态结构体
typedef struct {
    char filename[MAX_FILENAME_LENGTH];      // 文件名
    off_t file_size;                         // 文件大小
    int block_size;                          // 块大小
    int total_blocks;                        // 总块数
    int *downloaded_blocks;                  // 已下载块标记
    char temp_filename[MAX_FILENAME_LENGTH]; // 临时文件名
    char status_filename[MAX_FILENAME_LENGTH]; // 状态文件名
    char md5[33];                            // 文件MD5值
    FILE *temp_file;                         // 临时文件指针
    double progress;                         // 下载进度(0-100)
    double speed;                            // 下载速度(KB/s)
    time_t start_time;                       // 开始时间
    volatile int is_running;                 // 是否正在运行
} DownloadStatus;

DownloadStatus download_status;
int client_socket;

// 保存下载状态到文件
void save_download_status() {
    if (!download_status.is_running)
        return;

    FILE *status_file = fopen(download_status.status_filename, "w");
    if (!status_file) {
        perror("fopen status file");
        return;
    }

    // 写入基本信息
    fprintf(status_file, "filename: %s\n", download_status.filename);
    fprintf(status_file, "file_size: %ld\n", download_status.file_size);
    fprintf(status_file, "block_size: %d\n", download_status.block_size);
    fprintf(status_file, "total_blocks: %d\n", download_status.total_blocks);
    fprintf(status_file, "md5: %s\n", download_status.md5);

    // 写入已下载块信息
    fprintf(status_file, "downloaded_blocks:");
    for (int i = 0; i < download_status.total_blocks; i++) {
        if (download_status.downloaded_blocks[i]) {
            fprintf(status_file, " %d", i);
        }
    }
    fprintf(status_file, "\n");

    fclose(status_file);
    printf("\nDownload status saved to %s\n", download_status.status_filename);
}

// 加载下载状态从文件
int load_download_status(const char *filename) {
    // 构建状态文件名
    snprintf(download_status.status_filename, sizeof(download_status.status_filename),
             "%s%s", filename, STATUS_FILE_EXT);

    FILE *status_file = fopen(download_status.status_filename, "r");
    if (!status_file) {
        return -1; // 状态文件不存在
    }

    char line[256];
    char key[64];
    char value[256];

    // 初始化下载状态
    memset(&download_status, 0, sizeof(download_status));
    strncpy(download_status.filename, filename, MAX_FILENAME_LENGTH);

    while (fgets(line, sizeof(line), status_file)) {
        if (sscanf(line, "%63[^:]: %255[^\n]", key, value) == 2) {
            if (strcmp(key, "filename") == 0) {
                strncpy(download_status.filename, value, MAX_FILENAME_LENGTH);
            } else if (strcmp(key, "file_size") == 0) {
                download_status.file_size = atoll(value);
            } else if (strcmp(key, "block_size") == 0) {
                download_status.block_size = atoi(value);
            } else if (strcmp(key, "total_blocks") == 0) {
                download_status.total_blocks = atoi(value);
                download_status.downloaded_blocks = calloc(download_status.total_blocks, sizeof(int));
                if (!download_status.downloaded_blocks) {
                    perror("calloc");
                    fclose(status_file);
                    return -1;
                }
            } else if (strcmp(key, "md5") == 0) {
                strncpy(download_status.md5, value, 32);
            } else if (strcmp(key, "downloaded_blocks") == 0) {
                char *token = strtok(value, " ");
                while (token) {
                    int block_index = atoi(token);
                    if (block_index >= 0 && block_index < download_status.total_blocks) {
                        download_status.downloaded_blocks[block_index] = 1;
                    }
                    token = strtok(NULL, " ");
                }
            }
        }
    }

    fclose(status_file);

    // 构建临时文件名
    snprintf(download_status.temp_filename, sizeof(download_status.temp_filename),
             "%s%s", filename, TEMP_FILE_EXT);

    // 检查临时文件是否存在
    struct stat temp_file_stat;
    if (stat(download_status.temp_filename, &temp_file_stat) == -1) {
        printf("Temporary file %s not found, cannot resume download\n", download_status.temp_filename);
        free(download_status.downloaded_blocks);
        return -1;
    }

    // 打开临时文件
    download_status.temp_file = fopen(download_status.temp_filename, "r+b");
    if (!download_status.temp_file) {
        perror("fopen temp file");
        free(download_status.downloaded_blocks);
        return -1;
    }

    // 计算已下载进度
    int downloaded_count = 0;
    for (int i = 0; i < download_status.total_blocks; i++) {
        if (download_status.downloaded_blocks[i]) {
            downloaded_count++;
        }
    }
    download_status.progress = (double)downloaded_count / download_status.total_blocks * 100;

    printf("Resuming download from %.1f%%\n", download_status.progress);
    return 0;
}

// 初始化新的下载
void init_new_download(const char *filename, const FileInfo *file_info) {
    memset(&download_status, 0, sizeof(download_status));
    strncpy(download_status.filename, filename, MAX_FILENAME_LENGTH);
    download_status.file_size = file_info->file_size;
    download_status.block_size = file_info->block_size;
    download_status.total_blocks = file_info->total_blocks;
    strncpy(download_status.md5, file_info->md5, 32);

    // 分配已下载块标记数组
    download_status.downloaded_blocks = calloc(download_status.total_blocks, sizeof(int));
    if (!download_status.downloaded_blocks) {
        error_exit("calloc");
    }

    // 构建临时文件名和状态文件名
    snprintf(download_status.temp_filename, sizeof(download_status.temp_filename),
             "%s%s", filename, TEMP_FILE_EXT);
    snprintf(download_status.status_filename, sizeof(download_status.status_filename),
             "%s%s", filename, STATUS_FILE_EXT);

    // 创建临时文件
    download_status.temp_file = fopen(download_status.temp_filename, "wb");
    if (!download_status.temp_file) {
        error_exit("fopen temp file");
    }

    // 设置文件大小（创建稀疏文件，如果系统支持）
    #ifdef __linux__
    if (ftruncate(fileno(download_status.temp_file), download_status.file_size) == -1) {
        perror("ftruncate");
        // 忽略错误，继续执行
    }
    #endif

    download_status.start_time = time(NULL);
    download_status.is_running = 1;
}

// 信号处理函数
void signal_handler(int signum) {
    if (signum == SIGINT) {
        printf("\nReceived Ctrl+C, stopping download...\n");
        download_status.is_running = 0;
        save_download_status();
        
        if (download_status.temp_file) {
            fclose(download_status.temp_file);
        }
        
        close(client_socket);
        exit(EXIT_SUCCESS);
    }
}

// 检查是否可以续传
int check_resume(const char *filename) {
    char status_filename[MAX_FILENAME_LENGTH];
    snprintf(status_filename, sizeof(status_filename), "%s%s", filename, STATUS_FILE_EXT);

    struct stat status_file_stat;
    if (stat(status_filename, &status_file_stat) == -1) {
        return 0; // 状态文件不存在，不能续传
    }

    // 加载下载状态
    if (load_download_status(filename) == 0) {
        return 1; // 可以续传
    } else {
        // 加载失败，删除状态文件
        unlink(status_filename);
        return 0;
    }
}

// 请求文件信息
int request_file_info(const char *filename, FileInfo *file_info) {
    // 发送文件请求
    if (send_message(client_socket, MSG_FILE_REQUEST, filename, strlen(filename) + 1) < 0) {
        return -1;
    }

    // 接收响应
    MessageHeader header;
    void *data = NULL;
    ssize_t recv_len = receive_message(client_socket, &header, &data);
    if (recv_len <= 0) {
        return -1;
    }

    if (header.type == MSG_FILE_INFO) {
        // 接收文件信息
        if (header.length == sizeof(FileInfo)) {
            memcpy(file_info, data, sizeof(FileInfo));
            free(data);
            return 0;
        } else {
            printf("Invalid file info length: %d\n", header.length);
            free(data);
            return -1;
        }
    } else if (header.type == MSG_TRANSFER_STATUS) {
        // 接收错误信息
        TransferStatusMsg *status_msg = (TransferStatusMsg *)data;
        printf("Error: %s\n", status_msg->message);
        free(data);
        return -1;
    } else {
        printf("Unexpected message type: %d\n", header.type);
        free(data);
        return -1;
    }
}

// 请求块数据
int request_block(int block_index) {
    BlockRequest request;
    strncpy(request.filename, download_status.filename, MAX_FILENAME_LENGTH);
    request.block_index = block_index;

    // 发送块请求
    if (send_message(client_socket, MSG_BLOCK_REQUEST, &request, sizeof(BlockRequest)) < 0) {
        return -1;
    }

    // 接收响应
    MessageHeader header;
    void *data = NULL;
    ssize_t recv_len = receive_message(client_socket, &header, &data);
    if (recv_len <= 0) {
        return -1;
    }

    if (header.type == MSG_BLOCK_DATA) {
        // 接收块数据
        BlockData *block_data = (BlockData *)data;
        
        // 验证块信息
        if (strcmp(block_data->filename, download_status.filename) != 0 ||
            block_data->block_index != block_index) {
            printf("Received invalid block: expected %s:%d, got %s:%d\n",
                   download_status.filename, block_index,
                   block_data->filename, block_data->block_index);
            free(data);
            return -1;
        }

        // 验证块MD5
        char calculated_md5[33];
        calculate_block_md5(block_data->data, block_data->block_size, calculated_md5);
        if (strcmp(block_data->md5, calculated_md5) != 0) {
            printf("Block %d MD5 mismatch: expected %s, got %s\n",
                   block_index, calculated_md5, block_data->md5);
            free(data);
            return -1;
        }

        // 定位到块位置
        off_t offset = (off_t)block_index * download_status.block_size;
        if (fseek(download_status.temp_file, offset, SEEK_SET) != 0) {
            perror("fseek");
            free(data);
            return -1;
        }

        // 写入块数据
        size_t bytes_written = fwrite(block_data->data, 1, block_data->block_size, download_status.temp_file);
        if (bytes_written != block_data->block_size) {
            perror("fwrite");
            free(data);
            return -1;
        }

        // 标记块为已下载
        download_status.downloaded_blocks[block_index] = 1;

        // 更新进度
        int downloaded_count = 0;
        for (int i = 0; i < download_status.total_blocks; i++) {
            if (download_status.downloaded_blocks[i]) {
                downloaded_count++;
            }
        }
        download_status.progress = (double)downloaded_count / download_status.total_blocks * 100;

        // 计算下载速度
        time_t current_time = time(NULL);
        double elapsed = difftime(current_time, download_status.start_time);
        if (elapsed > 0) {
            off_t downloaded_size = downloaded_count * download_status.block_size;
            if (downloaded_size > download_status.file_size) {
                downloaded_size = download_status.file_size;
            }
            download_status.speed = (double)downloaded_size / (1024 * elapsed);
        }

        // 显示进度
        print_progress(download_status.filename, download_status.progress,
                      download_status.speed, current_time - download_status.start_time);

        free(data);
        return 0;
    } else if (header.type == MSG_TRANSFER_STATUS) {
        // 接收错误信息
        TransferStatusMsg *status_msg = (TransferStatusMsg *)data;
        printf("\nError: %s\n", status_msg->message);
        free(data);
        return -1;
    } else {
        printf("\nUnexpected message type: %d\n", header.type);
        free(data);
        return -1;
    }
}

// 完成下载
void complete_download() {
    // 关闭临时文件
    fclose(download_status.temp_file);

    // 验证文件MD5
    char calculated_md5[33];
    if (calculate_file_md5(download_status.temp_filename, calculated_md5) == 0) {
        if (strcmp(calculated_md5, download_status.md5) != 0) {
            printf("\nFile MD5 mismatch: expected %s, got %s\n",
                   download_status.md5, calculated_md5);
            printf("Downloaded file may be corrupted\n");
        } else {
            printf("\nFile MD5 verified successfully\n");
        }
    } else {
        printf("\nFailed to calculate MD5 for downloaded file\n");
    }

    // 重命名临时文件为正式文件
    if (rename(download_status.temp_filename, download_status.filename) == 0) {
        printf("\nDownload completed successfully: %s\n", download_status.filename);
    } else {
        perror("\nrename");
        printf("Downloaded file saved as: %s\n", download_status.temp_filename);
    }

    // 删除状态文件
    unlink(download_status.status_filename);

    // 释放资源
    free(download_status.downloaded_blocks);
}

int main(int argc, char *argv[]) {
    if (argc != 4) {
        fprintf(stderr, "Usage: %s <server_ip> <port> <filename>\n", argv[0]);
        exit(EXIT_FAILURE);
    }

    const char *server_ip = argv[1];
    int port = atoi(argv[2]);
    const char *filename = argv[3];

    // 注册信号处理函数
    signal(SIGINT, signal_handler);

    // 创建套接字
    client_socket = socket(AF_INET, SOCK_STREAM, 0);
    if (client_socket < 0) {
        error_exit("socket");
    }

    // 设置服务器地址
    struct sockaddr_in server_addr;
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(port);

    // 转换IP地址
    if (inet_pton(AF_INET, server_ip, &server_addr.sin_addr) <= 0) {
        error_exit("inet_pton");
    }

    // 连接服务器
    printf("Connecting to %s:%d...\n", server_ip, port);
    if (connect(client_socket, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        error_exit("connect");
    }
    printf("Connected successfully\n");

    // 检查是否可以续传
    int can_resume = check_resume(filename);
    FileInfo file_info;

    if (can_resume) {
        // 请求文件信息以验证
        if (request_file_info(filename, &file_info) < 0) {
            printf("Failed to request file info, cannot resume download\n");
            can_resume = 0;
        } else {
            // 验证文件信息是否匹配
            if (file_info.file_size != download_status.file_size ||
                strcmp(file_info.md5, download_status.md5) != 0) {
                printf("File information mismatch, cannot resume download\n");
                printf("Starting new download...\n");
                can_resume = 0;
                
                // 关闭临时文件并释放资源
                fclose(download_status.temp_file);
                free(download_status.downloaded_blocks);
            }
        }
    }

    if (!can_resume) {
        // 请求文件信息
        printf("Requesting file info for %s...\n", filename);
        if (request_file_info(filename, &file_info) < 0) {
            close(client_socket);
            exit(EXIT_FAILURE);
        }

        // 初始化新的下载
        init_new_download(filename, &file_info);
        printf("Starting new download: %s\n", filename);
        printf("File size: %ld bytes, Blocks: %d, Block size: %d bytes\n",
               file_info.file_size, file_info.total_blocks, file_info.block_size);
    }

    download_status.is_running = 1;
    download_status.start_time = time(NULL);

    // 下载未完成的块
    while (download_status.is_running) {
        int all_downloaded = 1;
        int blocks_downloaded = 0;

        for (int i = 0; i < download_status.total_blocks && download_status.is_running; i++) {
            if (!download_status.downloaded_blocks[i]) {
                all_downloaded = 0;
                
                // 请求块数据
                if (request_block(i) < 0) {
                    printf("\nFailed to download block %d, retrying...\n", i);
                    sleep(1); // 等待1秒后重试
                } else {
                    blocks_downloaded++;
                    
                    // 每下载10个块保存一次状态
                    if (blocks_downloaded % 10 == 0) {
                        save_download_status();
                    }
                }
            }
        }

        if (all_downloaded) {
            break;
        }
    }

    if (download_status.is_running) {
        // 完成下载
        complete_download();
    } else {
        // 保存下载状态
        save_download_status();
    }

    // 清理资源
    close(client_socket);
    return 0;
}
